
import streamlit as st
from core.utils.path import upload_dir
from core.novel_analyze import NovelAnalyzer
import os
import torch
torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)]


def main():

    topn = int(st.number_input('分析的章节数', min_value=1,
                               max_value=10000, value=100, step=1))
    max_len = int(st.number_input('模型输入字数限制', min_value=10000,
                                  max_value=128000, value=32000, step=1000))
    st.title(f'小说分析(前{topn}章，{max_len}字切片)')

    uploaded_file = st.file_uploader("请选择文本文件进行上传", type=["txt"])
    if uploaded_file is not None:
        file_bytes = uploaded_file.read()
        save_path = upload_dir / uploaded_file.name
        with open(save_path, "wb") as f:
            f.write(file_bytes)

        analyzer = NovelAnalyzer(save_path, topn=topn, max_len=max_len)
        novel = analyzer.analyze()

        st.markdown(f"**总得分：{novel.total_score}**")

        st.markdown("**情节设计评分**")
        st.json(novel.plot_rating, expanded=True)

        st.markdown("**人物塑造评分**")
        st.json(novel.character_rating, expanded=True)

        st.markdown("**世界观构建评分**")
        st.json(novel.worldview_rating, expanded=True)

        st.markdown("**文笔写作评分**")
        st.json(novel.style_rating, expanded=True)

        st.markdown("**各章节的问题点**")
        st.json({chapter.chapter_title: chapter.problems for chapter in novel.chapters})

        st.markdown("**各章节的情节设计分析**")
        st.json({chapter.chapter_title: chapter.plot for chapter in novel.chapters})

        st.markdown("**各章节的人物塑造分析**")
        st.json({chapter.chapter_title: chapter.characterization for chapter in novel.chapters})

        st.markdown("**各章节的世界观构建分析**")
        st.json({chapter.chapter_title: chapter.worldview for chapter in novel.chapters})

        st.markdown("**各章节的文笔写作分析**")
        st.json({chapter.chapter_title: chapter.writing_style for chapter in novel.chapters})


if __name__ == "__main__":
    main()